#include "yolov8.h"

void query_dfl_len(MODEL_INFO *m, YOLO_INFO *y_info) {
  // set dfl_len
  if ((y_info->m_type == YOLOV8)) {
    if (m->n_output > 6) {
      y_info->score_sum_available = true;
    } else {
      y_info->score_sum_available = false;
    }

    if ((y_info->m_type == YOLOV8)) {
      y_info->dfl_len = (int)(m->out_attr[0].dims[2] / 4);
    }
  }
}

YoloV8::YoloV8(const std::string &rknnModelPath) {
  model_info_.m_path = const_cast<char *>(rknnModelPath.c_str());
  y_info_.m_type = YOLOV8;
  y_info_.anchor_per_branch = 1;
  y_info_.post_type = Q8;
  y_info_.in_source = SINGLE_IMG;

  printf("Loading mode : %s...\n", rknnModelPath.c_str());

  rkdemo_init(&model_info_);
  rkdemo_init_input_buffer_all(&model_info_, NORMAL_API, RKNN_TENSOR_UINT8);
  // for (int i=0; i< model_info_.n_input; i++) {
  //     rknn_set_io_mem(model_info_.ctx, model_info_.input_mem[i],
  //     &model_info_.in_attr[i]);
  // }

  if (y_info_.post_type == Q8) {
    rkdemo_init_output_buffer_all(&model_info_, NORMAL_API, 0);
  } else {
    rkdemo_init_output_buffer_all(&model_info_, NORMAL_API, 1);
  }

  // for (int i=0; i< model_info_.n_output; i++) {
  //   rknn_set_io_mem(model_info_.ctx, model_info_.output_mem[i],
  //   &model_info_.out_attr[i]);
  // }

  query_dfl_len(&model_info_, &y_info_);

  // rknn_run(model_info_.ctx, NULL);

  printf("load yolov8 model finish.");
}

int YoloV8::predict(const cv::Mat &input,
                    detect_result_group_t &detect_result_group) {
  int ret = -1;
  float scale_w, scale_h;
  // cv::Mat resize_img = preprocess(input, scale_w, scale_h);
  cv::Mat resize_img =
      cv::Mat::zeros(model_input_size_, model_input_size_, CV_8UC3);
  preprocess_with_rga(input, resize_img.data, scale_w, scale_h);
  // cv::imwrite("resize_img.jpg", resize_img);

  model_info_.inputs[0].buf = resize_img.data;
  rknn_inputs_set(model_info_.ctx, model_info_.n_input, model_info_.inputs);

  // memcpy(model_info_.input_mem[0]->logical_addr, resize_img.data,
  // model_info_.in_attr[0].size); printf("model_info_.in_attr[0].size : %d \n",
  // model_info_.in_attr[0].size); cv::Mat image1(640, 640, CV_8UC3,
  // (uint8_t*)model_info_.input_mem[0]->logical_addr);
  // cv::imwrite("image1.jpg", image1);

  ret = rknn_run(model_info_.ctx, NULL);

  ret = rknn_outputs_get(model_info_.ctx, model_info_.n_output,
                         model_info_.outputs, NULL);

  ret = postprocess(input.cols, input.rows, scale_w, scale_h,
                    detect_result_group);

  return 0;
}

int YoloV8::predict_with_nv12(MEDIA_BUFFER fram_mb, cv::Size frame_size,
                              detect_result_group_t &detect_result_group) {
  int ret = -1;
  float scale_w, scale_h;
  // cv::Mat resize_img = preprocess(input, scale_w, scale_h);
  cv::Mat resize_img =
      cv::Mat::zeros(model_input_size_, model_input_size_, CV_8UC3);
  preprocess_with_rga(fram_mb, frame_size, resize_img.data, scale_w, scale_h);
  // cv::imwrite("resize_img.jpg", resize_img);

  model_info_.inputs[0].buf = resize_img.data;
  rknn_inputs_set(model_info_.ctx, model_info_.n_input, model_info_.inputs);

  // memcpy(model_info_.input_mem[0]->logical_addr, resize_img.data,
  // model_info_.in_attr[0].size); printf("model_info_.in_attr[0].size : %d \n",
  // model_info_.in_attr[0].size); cv::Mat image1(640, 640, CV_8UC3,
  // (uint8_t*)model_info_.input_mem[0]->logical_addr);
  // cv::imwrite("image1.jpg", image1);

  ret = rknn_run(model_info_.ctx, NULL);

  ret = rknn_outputs_get(model_info_.ctx, model_info_.n_output,
                         model_info_.outputs, NULL);

  ret = postprocess(frame_size.width, frame_size.height, scale_w, scale_h,
                    detect_result_group);

  return 0;
}

void YoloV8::preprocess_with_rga(const cv::Mat &image, void *dst_buffer,
                                 float &scale_w, float &scale_h) {
  int col = image.cols;
  int row = image.rows;
  int _max = MAX(col, row);

  int pad_w = _max - col;
  int pad_h = _max - row;
  float scale = model_input_size_ * 1.f / _max;

  scale_w = (model_input_size_ - pad_w * scale) * 1.f / col;
  scale_h = (model_input_size_ - pad_h * scale) * 1.f / row;

  rga_buffer_t src = wrapbuffer_virtualaddr(
      image.data, col, row,
      RK_FORMAT_BGR_888);  // 需要4字节对齐，因为固定上层输入1280x720就不进行处理了
  rga_buffer_t dst =
      wrapbuffer_virtualaddr(dst_buffer, model_input_size_, model_input_size_,
                             RK_FORMAT_BGR_888);  // RGA对RGB888要4字节对齐
  rga_buffer_t pat;

  im_rect src_rect = {0, 0, col, row};
  int crop_w = int(model_input_size_ - pad_w * scale);
  int crop_h = int(model_input_size_ - pad_h * scale);

  im_rect crop_rect = {0, 0, crop_w, crop_h};
  im_rect prect;

  IM_STATUS STATUS =
      improcess(src, dst, pat, src_rect, crop_rect, prect, IM_SYNC);

  if (STATUS != IM_STATUS_SUCCESS) {
    printf("improcess failed: %s", imStrError(STATUS));
    return;
  }
}

void YoloV8::preprocess_with_rga(MEDIA_BUFFER fram_mb, cv::Size frame_size,
                                 void *dst_buffer, float &scale_w,
                                 float &scale_h) {
  int col = frame_size.width;
  int row = frame_size.height;
  int _max = MAX(col, row);

  int pad_w = _max - col;
  int pad_h = _max - row;
  float scale = model_input_size_ * 1.f / _max;

  scale_w = (model_input_size_ - pad_w * scale) * 1.f / col;
  scale_h = (model_input_size_ - pad_h * scale) * 1.f / row;

  rga_buffer_t src = wrapbuffer_fd(
      RK_MPI_MB_GetFD(fram_mb), col, row,
      RK_FORMAT_YCbCr_420_SP);  // 需要4字节对齐，因为固定上层输入1280x720就不进行处理了
  rga_buffer_t dst =
      wrapbuffer_virtualaddr(dst_buffer, model_input_size_, model_input_size_,
                             RK_FORMAT_BGR_888);  // RGA对RGB888要4字节对齐
  rga_buffer_t pat;

  im_rect src_rect = {0, 0, col, row};
  int crop_w = int(model_input_size_ - pad_w * scale);
  int crop_h = int(model_input_size_ - pad_h * scale);

  im_rect crop_rect = {0, 0, crop_w, crop_h};
  im_rect prect;

  IM_STATUS STATUS =
      improcess(src, dst, pat, src_rect, crop_rect, prect, IM_SYNC);

  if (STATUS != IM_STATUS_SUCCESS) {
    printf("improcess failed: %s", imStrError(STATUS));
    return;
  }
}

cv::Mat YoloV8::preprocess(const cv::Mat &image, float &scale_w,
                           float &scale_h) {
  int col = image.cols;
  int row = image.rows;
  int _max = MAX(col, row);
  cv::Mat result = cv::Mat::zeros(_max, _max, CV_8UC3);
  image.copyTo(result(cv::Rect(0, 0, col, row)));

  cv::resize(result, result, cv::Size(model_input_size_, model_input_size_));

  int pad_w = _max - col;
  int pad_h = _max - row;
  float scale = model_input_size_ * 1.f / _max;

  scale_w = (model_input_size_ - pad_w * scale) * 1.f / col;
  scale_h = (model_input_size_ - pad_h * scale) * 1.f / row;

  return result;
}

int YoloV8::postprocess(int input_w, int input_h, float scale_w, float scale_h,
                        detect_result_group_t &detect_result_group) {
  /* Post process */
  void *output_buf_list[model_info_.n_output];
  for (int i = 0; i < model_info_.n_output; i++) {
    // output_buf_list[i] = model_info_.output_mem[i]->logical_addr;
    output_buf_list[i] = model_info_.outputs[i].buf;
  }

  post_process(output_buf_list, &model_info_, &y_info_, &detect_result_group);

  for (int i = 0; i < detect_result_group.count; i++) {
    detect_result_group.results[i].box.left /= scale_w;
    detect_result_group.results[i].box.right /= scale_w;
    detect_result_group.results[i].box.top /= scale_h;
    detect_result_group.results[i].box.bottom /= scale_h;

    detect_result_group.results[i].box.left =
        std::max(0, std::min(input_w, detect_result_group.results[i].box.left));
    detect_result_group.results[i].box.right = std::max(
        0, std::min(input_w, detect_result_group.results[i].box.right));
    detect_result_group.results[i].box.top =
        std::max(0, std::min(input_h, detect_result_group.results[i].box.top));
    detect_result_group.results[i].box.bottom = std::max(
        0, std::min(input_h, detect_result_group.results[i].box.bottom));

    // detect_result_t *det_result = &detect_result_group.results[i];
    // printf("%s%s @ (%d %d %d %d) %f\n", "---------------",
    //             det_result->name,
    //             det_result->box.left, det_result->box.top,
    //             det_result->box.right, det_result->box.bottom,
    //             det_result->prop);
  }

  return 0;
}
